# -*- coding: utf-8 -*-
"""
LLM 交互服务 - 封装 LangChain + Ollama 集成

@author: AI Assistant
"""
import json
from typing import Optional, Dict, Any, List
from loguru import logger

from langchain_ollama import ChatOllama
from langchain.schema import SystemMessage, HumanMessage, AIMessage

import sys
from pathlib import Path
project_root = Path(__file__).resolve().parent.parent
if str(project_root) not in sys.path:
    sys.path.insert(0, str(project_root))

from config import OLLAMA_CONFIG, LLM_GENERATION_CONFIG
from models.domain_models import LLMResponse


class LLMService:
    """LLM 服务类 - 封装 Ollama 模型交互"""
    
    def __init__(self):
        """初始化 LLM 服务"""
        try:
            # 初始化 Ollama Chat模型
            self.llm = ChatOllama(
                model=OLLAMA_CONFIG["chat_model"],
                base_url=OLLAMA_CONFIG["base_url"],
                temperature=LLM_GENERATION_CONFIG["temperature"],
                num_predict=LLM_GENERATION_CONFIG["max_tokens"],
                top_p=LLM_GENERATION_CONFIG["top_p"],
                format="json"  # 强制返回 JSON 格式
            )
            logger.info(f"✅ LLM 服务初始化成功 - 模型: {OLLAMA_CONFIG['chat_model']}")
        except Exception as e:
            logger.error(f"❌ LLM 服务初始化失败: {str(e)}")
            raise
    
    def generate_response(
        self,
        system_prompt: str,
        user_message: str,
        conversation_history: Optional[List[Dict[str, str]]] = None
    ) -> Optional[Dict[str, Any]]:
        """生成 LLM 响应
        
        Args:
            system_prompt: 系统提示词
            user_message: 用户消息
            conversation_history: 对话历史，格式：[{"role": "user/assistant", "content": "..."}]
            
        Returns:
            解析后的 LLM 响应字典，如果失败则返回 None
        """
        try:
            # 构建消息列表
            messages = [SystemMessage(content=system_prompt)]
            
            # 添加对话历史（只保留最近几轮）
            if conversation_history:
                for msg in conversation_history[-5:]:  # 保留最近5轮对话
                    if msg["role"] == "user":
                        messages.append(HumanMessage(content=msg["content"]))
                    elif msg["role"] == "assistant":
                        messages.append(AIMessage(content=msg["content"]))
            
            # 添加当前用户消息
            messages.append(HumanMessage(content=user_message))
            
            logger.info(f"发送请求到 LLM - 消息数: {len(messages)}")
            
            # 调用 LLM
            response = self.llm.invoke(messages)
            response_content = response.content
            
            logger.info(f"LLM 原始响应长度: {len(response_content)} 字符")
            logger.debug(f"LLM 原始响应: {response_content[:1000]}...")
            
            # 解析 JSON 响应
            parsed_response = self._parse_llm_response(response_content)
            
            if parsed_response:
                logger.info(f"✅ LLM 响应解析成功 - understood: {parsed_response.get('understood')}")
                return parsed_response
            else:
                logger.warning("❌ LLM 响应解析失败")
                return self._create_fallback_response(user_message, response_content)
                
        except Exception as e:
            logger.error(f"LLM 生成响应失败: {str(e)}", exc_info=True)
            return self._create_error_response(str(e))
    
    def _parse_llm_response(self, response_content: str) -> Optional[Dict[str, Any]]:
        """解析 LLM 的 JSON 响应
        
        Args:
            response_content: LLM 返回的原始内容
            
        Returns:
            解析后的字典，如果解析失败则返回 None
        """
        try:
            # 尝试直接解析 JSON
            response_dict = json.loads(response_content)
            
            # 验证必需字段
            if "understood" in response_dict and "message" in response_dict:
                return response_dict
            else:
                logger.warning("LLM 响应缺少必需字段")
                return None
                
        except json.JSONDecodeError:
            # 如果不是纯 JSON，尝试提取 JSON 代码块
            logger.warning("LLM 响应不是纯 JSON，尝试提取 JSON 代码块")
            
            # 查找 ```json ... ``` 代码块
            if "```json" in response_content:
                try:
                    start = response_content.find("```json") + 7
                    end = response_content.find("```", start)
                    json_str = response_content[start:end].strip()
                    response_dict = json.loads(json_str)
                    
                    if "understood" in response_dict and "message" in response_dict:
                        return response_dict
                except Exception as e:
                    logger.error(f"提取 JSON 代码块失败: {e}")
            
            # 查找 ``` ... ``` 代码块（不带 json 标记）
            elif "```" in response_content:
                try:
                    start = response_content.find("```") + 3
                    end = response_content.find("```", start)
                    json_str = response_content[start:end].strip()
                    response_dict = json.loads(json_str)
                    
                    if "understood" in response_dict and "message" in response_dict:
                        return response_dict
                except Exception as e:
                    logger.error(f"提取代码块失败: {e}")
            
            # 尝试查找 { ... } JSON对象
            try:
                start_idx = response_content.find('{')
                end_idx = response_content.rfind('}')
                if start_idx != -1 and end_idx != -1 and end_idx > start_idx:
                    json_str = response_content[start_idx:end_idx+1]
                    response_dict = json.loads(json_str)
                    
                    if "understood" in response_dict and "message" in response_dict:
                        logger.info("成功从响应中提取JSON对象")
                        return response_dict
            except Exception as e:
                logger.error(f"提取JSON对象失败: {e}")
            
            return None
    
    def _create_fallback_response(
        self, 
        user_message: str, 
        raw_response: str
    ) -> Dict[str, Any]:
        """创建后备响应（当 LLM 未返回标准 JSON 时）
        
        Args:
            user_message: 用户消息
            raw_response: LLM 原始响应
            
        Returns:
            标准化的响应字典
        """
        return {
            "understood": False,
            "message": raw_response[:500],  # 截取前500个字符
            "parameter_changes": None,
            "explanation": "抱歉，我未能理解您的请求。请尝试更具体地描述您想要修改的参数。"
        }
    
    def _create_error_response(self, error_message: str) -> Dict[str, Any]:
        """创建错误响应
        
        Args:
            error_message: 错误消息
            
        Returns:
            错误响应字典
        """
        return {
            "understood": False,
            "message": f"抱歉，处理您的请求时发生错误：{error_message}",
            "parameter_changes": None,
            "explanation": "请稍后重试或联系管理员。"
        }
    
    def test_connection(self) -> bool:
        """测试 LLM 连接是否正常
        
        Returns:
            bool: 连接正常返回 True
        """
        try:
            test_message = [HumanMessage(content="你好")]
            response = self.llm.invoke(test_message)
            logger.info(f"✅ LLM 连接测试成功 - 响应: {response.content[:50]}...")
            return True
        except Exception as e:
            logger.error(f"❌ LLM 连接测试失败: {str(e)}")
            return False


# 使用示例
if __name__ == "__main__":
    logger.info("测试 LLM 服务...")
    
    try:
        # 初始化服务
        llm_service = LLMService()
        
        # 测试连接
        if llm_service.test_connection():
            logger.info("LLM 连接测试通过")
            
            # 测试对话
            system_prompt = """你是一个水文模型参数调整助手。
用户会告诉你要修改什么参数，你需要以 JSON 格式回复。

JSON 格式：
{
    "understood": true/false,
    "message": "你的回复",
    "parameter_changes": {"参数路径": 新值},
    "explanation": "说明"
}
"""
            
            user_message = "请将初始水位设置为 100 米"
            
            response = llm_service.generate_response(
                system_prompt=system_prompt,
                user_message=user_message
            )
            
            if response:
                logger.info(f"✅ 测试成功")
                logger.info(f"  - understood: {response.get('understood')}")
                logger.info(f"  - message: {response.get('message')}")
                logger.info(f"  - parameter_changes: {response.get('parameter_changes')}")
            else:
                logger.error("❌ 测试失败：未获取到响应")
                
    except Exception as e:
        logger.error(f"❌ 测试失败: {e}", exc_info=True)

